//
// Created by 郝羽 on 2018/3/21.
//
// srtp-decoder decrypt a pcap file using srtp

#include <getopt.h>
#include <sstream>
#include <string>
#include <pcap/pcap.h>

#include "decoder.hpp"
#include "packet.hpp"

std::map<std::string, Decoder*> decoders;
int total = 0;
int decoded = 0;
pcap_t* pcap = nullptr;
pcap_dumper_t* pcapDump = nullptr;
int linktype = 0;

void printHelp() {
    fprintf(stderr, "srtp-decoder -d <src ip>:<src port>=}<dst ip>:<dst port>,<key/salt> <input file> <output file>\n");
}

void handleDecoderArg(const char* arg) {
    fprintf(stderr, "arg=%s\n", arg);
    std::string hintAndKey(arg);
    auto sep = hintAndKey.find(',');
    if (sep == std::string::npos) {
        fprintf(stderr, "invalid -d arg");
        return;
    }
    auto hint = hintAndKey.substr(0, sep);
    auto key = hintAndKey.substr(sep + 1);
    fprintf(stderr, "hint %s key/salt %s\n", hint.c_str(), key.c_str());
    auto d = new Decoder(key);
    decoders.emplace(hint, d);
}

void handlePcap(u_char*, const struct pcap_pkthdr *h, const u_char *bytes) {
    auto pkt = Packet::Parse(h, bytes, linktype);
    if (!pkt) {
        return;
    }
    std::ostringstream fmt;
    fmt << std::string(inet_ntoa(pkt->src_.sin_addr)) << ":" << htons(pkt->src_.sin_port)
        << "=}"
        << std::string(inet_ntoa(pkt->dst_.sin_addr)) << ":" << htons(pkt->dst_.sin_port);
    auto hint = fmt.str();
    // fprintf(stderr, "fmt=%s\n", hint.c_str());
    auto it = decoders.find(hint);
    if (it == decoders.end()) {
        return;
    }
    auto d = it->second;
    if (d->Decode(pkt)) {
        decoded++;    
        pcap_dump((u_char*)pcapDump, &pkt->pcapHeader_, pkt->capturedData_);
    } else {
        fprintf(stderr, "failed to decode one packets\n");
    }
}

int main(int argc, char* argv[]) {
    srtp_init();
    char pcapErrBuf[PCAP_ERRBUF_SIZE] = {0,};
    int c;
    while((c = getopt(argc, argv, "d:h")) != -1) {
        switch (c) {
            case 'd':
                handleDecoderArg(optarg);
                break;
            case 'h':
                printHelp();
                return 0;
            case '?':
                if (optopt == 'd'){
                    fprintf (stderr, "option -%c requires an argument.\n", optopt);
                }
                printHelp();
                return 1;
            default:
                fprintf (stderr, "unknown option `-%c'.\n", optopt);
                printHelp();
                return 1;
        }
    }

    std::string fin("-");
    std::string fout("-");

    if (optind > argc - 2) {
        fprintf(stderr, "input file not provided, using stdin\n");
        fin = "-";
    } else {
        fprintf(stderr, "input file is %s", argv[optind]);
        fin = std::string(argv[optind]);
        optind++;
        if (optind < argc - 2) {
            fprintf(stderr, "output file not provided, using stdout\n");
            fout = "-";
        } else {
            fprintf(stderr, "output file is %s", argv[optind]);
            fout = std::string(argv[optind]);
            optind++;
        }
    }

    pcap = pcap_open_offline(fin.c_str(), pcapErrBuf);
    if (!pcap) {
        fprintf(stderr, "could not read input: %s\n", pcapErrBuf);
        return 10;
    }

    linktype = pcap_datalink(pcap);
    switch(linktype) {
        case DLT_EN10MB:
        case DLT_NULL:
        case DLT_RAW:
            break;
        default:
            fprintf(stderr, " unsupported data link type: %s: %s\n", pcap_datalink_val_to_name(linktype), pcap_datalink_val_to_description(linktype));
            return 20;

    }
    fprintf(stderr, "data link is %s: %s\n", pcap_datalink_val_to_name(linktype), pcap_datalink_val_to_description(linktype));

    pcapDump = pcap_dump_open(pcap, fout.c_str());
    if (!pcapDump) {
        fprintf(stderr, "could not write output: %s\n", pcap_geterr(pcap));
        return 30;
    }

    auto res = pcap_loop(pcap, -1, handlePcap, nullptr);
    if (res == -1) {
        fprintf(stderr, "%s", pcap_geterr(pcap));
        return 40;
    }

    fprintf(stderr, "done\n");
    
    fprintf(stderr, "decoded %d packets\n", decoded);
    pcap_dump_flush(pcapDump);
    pcap_dump_close(pcapDump);
    pcap_close(pcap);
}

